#include <src/pcps.h>
#include <src/input.h>
#include <src/hamiltonian.h>
#include <src/wavefunction.h>

//--------------------------------------------------------------------------------
// pcps::projected_driver -- performs a projected cps optimization of the
//                           wavefunction
//--------------------------------------------------------------------------------

template <class SCALAR>
void pcps::projected_driver(const pcps::Input & userinp, const pcps::Hamiltonian<SCALAR> & ham, pcps::Wavefunction<SCALAR> & wfn) {

  // get MPI info
  const MPI::Comm & comm = MPI::COMM_WORLD;
  const int nproc = comm.Get_size();
  const int myrank = comm.Get_rank();

  // set limits of macro iteration loop
  const int maxloop = userinp.max_macro_loop();
  const double tol = userinp.energy_tolerance();

  // start a timer
  pcps::Stopwatch sw; sw.start(); 

  // print a greeting
  if (myrank == 0) {
    std::cout << "  --- Optimizing the Wavefunction Projectively ---" << std::endl;
    std::cout << std::endl;
  }

  // set the perturbation level to zero for optimizing non-perturbative correlators
  int perturbation_level = 0;

  // print the initial energy
  SCALAR energy = pcps::energy(userinp, ham, wfn, perturbation_level);
  if (myrank == 0) {
    std::cout << boost::format("initial projected energy = %s") % pcps::format_number("%20.12f", energy) << std::endl;
    std::cout << std::endl;
  }

  // iteratively optimize the correlators
  bool is_converged = false;
  for (int iter = 0; iter < maxloop && !is_converged; iter++) {

    // start a timer for the macro iteration
    pcps::Stopwatch sw_mac; sw_mac.start(); 

    // record the energy of the previous iteration
    SCALAR last_energy = energy;

    // determine the damping to apply to the correlator update
    const double damping = userinp.damping() * std::pow(double(0.2), iter);

    // do a Newton-Raphson update for the first so many iterations
    if (iter < userinp.max_nr_iter()) {

      pcps::projected_nr_update(userinp, ham, wfn, perturbation_level, damping);

      // compute and print the new energy
      energy = pcps::energy(userinp, ham, wfn, perturbation_level);
      if (myrank == 0) {
        std::cout << boost::format("macro iter %3i    energy = %s") % iter % pcps::format_number("%20.12f", energy) << std::endl;
        std::cout << std::endl;
      }
      
      //// for the first few iterations, round off the correlator values to preserve symmetry
      //if (iter < 3)
      //  for (int c_ind = 0; c_ind < wfn.correlators().size(); c_ind++)
      //    if (!wfn.correlators().at(c_ind).fixed())
      //      for (int cmpd = 0; cmpd < wfn.correlators().at(c_ind).nelements(); cmpd++)
      //        wfn.correlators().at(c_ind)[cmpd] = pcps::round(wfn.correlators().at(c_ind)[cmpd], 3);

    // do a sweep update for the remaining iterations
    } else {

      // update each non-perturbative correlator
      for (int c_ind = 0; c_ind < wfn.correlators().size(); c_ind++) {

        // create a reference to the correlator we are optimizing
        pcps::Correlator<SCALAR> & corr = wfn.correlators().at(c_ind);

        // do not optimize fixed correlators
        if (corr.fixed()) continue;

        // do not optimize images of translationally invariant correlators
        if (corr.ti()) continue;

        // print which correlator is being updated
        if (myrank == 0) {
          std::cout << boost::format("updating correlator (");
          for (std::vector<int>::const_iterator site = corr.site_vec().begin(); site != corr.site_vec().end(); site++)
            std::cout << boost::format(" %2i") % *site;
          std::cout << boost::format("  )") << std::endl;
        }

        // update the correlator
        const double noise_magnitude = 0.0; //(iter < 1 ? 0.1 : 0.0);
        pcps::Stopwatch sw_update;
        sw_update.start();
        pcps::update_correlator(userinp, ham, wfn, corr, noise_magnitude, damping, perturbation_level);
        sw_update.stop();

        // compute and print the new energy
        pcps::Stopwatch sw_energy;
        sw_energy.start();
        energy = pcps::energy(userinp, ham, wfn, perturbation_level);
        sw_energy.stop();
        if (myrank == 0) {
          std::cout << boost::format("update took %14.4f seconds") % sw_update.elapsed_seconds() << std::endl;
          std::cout << boost::format("energy took %14.4f seconds") % sw_energy.elapsed_seconds() << std::endl;
          std::cout << boost::format("new projected energy = %s") % pcps::format_number("%20.12f", energy) << std::endl;
          std::cout << std::endl;
        }

        // print the updated correlator
        if (myrank == 0 && userinp.verbose_print()) {
          pcps::print_correlator(corr);
          std::cout << std::endl;
        }

      }

      // compute and print the new energy
      energy = pcps::energy(userinp, ham, wfn, perturbation_level);
      if (myrank == 0) {
        std::cout << boost::format("macro iter %3i    energy = %s") % iter % pcps::format_number("%20.12f", energy) << std::endl;
        std::cout << std::endl;
      }

    }

    // write correlators to a binary file
    if (myrank == 0)
      wfn.dump_binary(userinp);

    // check for convergence
    is_converged = ( std::abs(energy - last_energy) < tol );
    comm.Bcast((void*)&is_converged, sizeof(bool), MPI::CHAR, 0);

    // print the elapsed time
    sw_mac.stop();
    if (myrank == 0) {
      std::cout << boost::format("time for macro iteration: %14.4f sec") % sw_mac.elapsed_seconds() << std::endl;
      std::cout << std::endl;
    }

  }

  // print the result of the iteration
  if (myrank == 0) {
    if (is_converged)
      std::cout << std::endl
                << boost::format("pcps macro iterations converged!") << std::endl
                << std::endl;
    else
      std::cout << std::endl
                << boost::format("no convergence...") << std::endl
                << std::endl;
  }

  //// write correlators to a binary file for use in jesse's code
  //if (myrank == 0)
  //  wfn.dump_for_jesse(userinp, "dump_for_jesse.bin");

  // compute the 2nd order energy contribution from perturbative long-range pair correlators
  if (userinp.do_perturb()) {

    // iteratively optimize the perturbative correlators
    perturbation_level = 0;
    is_converged = false;
    for (int iter = 0; iter < 1 && !is_converged; iter++) {

      // start a timer for the macro iteration
      pcps::Stopwatch sw_mac; sw_mac.start(); 

      // update each correlator
      for (typename std::vector< pcps::Correlator<SCALAR> >::iterator corr = wfn.pert_corrs().begin(); corr != wfn.pert_corrs().end(); corr++) {

        // print which correlator is being updated
        if (myrank == 0) {
          std::cout << boost::format("updating perturbative correlator (");
          for (std::vector<int>::const_iterator site = corr->site_vec().begin(); site != corr->site_vec().end(); site++)
            std::cout << boost::format(" %2i") % *site;
          std::cout << boost::format("  )") << std::endl;
          std::cout << std::endl;
        }

        // update the correlator
        pcps::Stopwatch sw_update;
        sw_update.start();
        pcps::update_correlator(userinp, ham, wfn, *corr, 0.0, 0.0, perturbation_level);
        sw_update.stop();

        // subtract 1 from the correlator so that only the perturbative part is stored
        for (int i = 0; i < corr->nelements(); i++)
          (*corr)[i] -= 1.0;

        // compute and print the new energy
        //pcps::Stopwatch sw_energy;
        //sw_energy.start();
        //energy = pcps::energy(userinp, ham, wfn, perturbation_level);
        //sw_energy.stop();
        if (myrank == 0) {
          std::cout << boost::format("update took %14.4f seconds") % sw_update.elapsed_seconds() << std::endl;
          //std::cout << boost::format("energy took %14.4f seconds") % sw_energy.elapsed_seconds() << std::endl;
          //std::cout << boost::format("new energy = %20.12f") % energy << std::endl;
          std::cout << std::endl;
        }

        // print the updated correlator
        if (myrank == 0) {
          pcps::print_correlator(*corr);
          std::cout << std::endl;
        }

      }

      //// compute and print the new energy
      //double last_energy = macro_energy;
      //energy = pcps::energy(userinp, ham, wfn, perturbation_level);
      //macro_energy = energy;
      //if (myrank == 0) {
      //  std::cout << boost::format("macro iter %3i    energy = %.12f") % iter % energy << std::endl;
      //  std::cout << std::endl;
      //}

      //// convert the perturbative correlators to non-perturbative correlators and compute the "exact" energy
      //const int n_reg_corrs = wfn.correlators().size();
      //for (typename std::vector< pcps::Correlator<SCALAR> >::iterator corr = wfn.pert_corrs().begin(); corr != wfn.pert_corrs().end(); ) {
      //  for (int i = 0; i < corr->nelements(); i++)
      //    (*corr)[i] += 1.0;
      //  wfn.correlators().push_back(*corr);
      //  corr = wfn.pert_corrs().erase(corr);
      //}
      //perturbation_level = 0;
      //pcps::Stopwatch sw_energy;
      //sw_energy.start();
      //energy = pcps::energy(userinp, ham, wfn, perturbation_level);
      //sw_energy.stop();
      //if (myrank == 0) {
      //  std::cout << boost::format("energy with non-perturbative long range took %14.4f seconds") % sw_energy.elapsed_seconds() << std::endl;
      //  std::cout << boost::format("energy with non-perturbative long range =    %.12f") % energy << std::endl;
      //  std::cout << std::endl;
      //}
      //perturbation_level = 2;
      //for (typename std::vector< pcps::Correlator<SCALAR> >::iterator corr = wfn.correlators().begin() + n_reg_corrs; corr != wfn.correlators().end(); ) {
      //  for (int i = 0; i < corr->nelements(); i++)
      //    (*corr)[i] -= 1.0;
      //  wfn.pert_corrs().push_back(*corr);
      //  corr = wfn.correlators().erase(corr);
      //}

      //// print the elapsed time
      //if (myrank == 0) {
      //  std::cout << boost::format("time for macro iteration: %14.4f sec") % sw_mac.elapsed_seconds() << std::endl;
      //  std::cout << std::endl;
      //}

      //// check for convergence
      //is_converged = ( std::fabs(macro_energy - last_energy) < tol );

    }

    // print the energy with 2nd order perturbative correlator corrections
    perturbation_level = 2;
    pcps::Stopwatch sw_energy;
    sw_energy.start();
    energy = pcps::energy(userinp, ham, wfn, perturbation_level);
    sw_energy.stop();
    if (myrank == 0) {
      std::cout << boost::format("2nd order perturbative projected energy took %14.4f seconds") % sw_energy.elapsed_seconds() << std::endl;
      std::cout << boost::format("2nd order perturbative projected energy =    %.12f") % energy << std::endl;
      std::cout << std::endl;
    }

  }

  // print the elapsed time
  if (myrank == 0) {
    std::cout << boost::format("time to optimize wavefunction: %.2f sec") % sw.elapsed_seconds() << std::endl;
    std::cout << std::endl;
  }

}

// explicitly instantiate the template
template void pcps::projected_driver(const pcps::Input &, const pcps::Hamiltonian< double               > &, pcps::Wavefunction< double               > &);
template void pcps::projected_driver(const pcps::Input &, const pcps::Hamiltonian< std::complex<double> > &, pcps::Wavefunction< std::complex<double> > &);
